import pandas as pd
import numpy as np
import torch
import torch.optim as optim
from data_processing import load_data, preprocess_data, prepare_data
from model import create_model

def train_model(model, X_train, y_train, epochs=50, batch_size=32):
    criterion = torch.nn.MSELoss()
    optimizer = optim.Adam(model.parameters())
    
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for i in range(0, len(X_train), batch_size):
            batch_X = X_train[i:i+batch_size]
            batch_y = y_train[i:i+batch_size]
            
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Loss: {total_loss/len(X_train):.4f}')

def main():
    # 加载并预处理训练数据
    print("正在加载并预处理训练数据...")
    train_df = load_data('train.csv')
    train_df = preprocess_data(train_df)

    # 准备模型输入数据
    print("正在准备模型输入数据...")
    X_train, X_test, y_train, y_test, scaler = prepare_data(train_df)

    # 创建并训练模型
    print("正在创建并训练模型...")
    model = create_model((X_train.shape[1], X_train.shape[2]))
    train_model(model, X_train, y_train)

    # 加载测试数据
    print("正在加载并预处理测试数据...")
    test_df = load_data('testA.csv')
    test_df = preprocess_data(test_df)

    # 准备测试数据
    print("正在准备测试数据以进行预测...")
    test_features = ['lat', 'lon', 'hour', 'day', 'month']
    test_scaled = scaler.transform(test_df[test_features])
    X_test_seq = torch.FloatTensor([test_scaled[i-10:i] for i in range(10, len(test_scaled))])

    # 进行预测
    print("正在进行预测...")
    model.eval()
    with torch.no_grad():
        predictions = model(X_test_seq).numpy()
    
    # 创建一个新的全零数组，形状与原始特征匹配
    full_predictions = np.zeros((predictions.shape[0], len(test_features)))
    full_predictions[:, :2] = predictions  # 只填充前两列(lat, lon)
    
    # 反向转换
    predictions = scaler.inverse_transform(full_predictions)[:, :2]  # 只取前两列

    # 准备提交文件
    print("正在准备提交文件...")
    submission = pd.DataFrame({
        'mmsi': test_df['mmsi'].iloc[10:].values,
        'lat': predictions[:, 0],
        'lon': predictions[:, 1]
    })

    # 保存提交文件
    submission.to_csv('submission.csv', index=False)
    print("提交文件已成功创建！")

if __name__ == "__main__":
    main()